#include <user/ulib.h>

physaddr_t
va2pa(void *addr)
{
	 if ((vpd[PAGEDIR_IDX(addr)] & PTE_P) &&
		 (vpt[PAGE_NUM(addr)] & PTE_P))
		  return PTE_ADDR(vpt[PAGE_NUM(addr)]) | PAGE_OFF(addr);
	 else return 0;
}

/* heap module ------------------------------------------------------------ */
static void *malloc_unsafe(size_t size);
static void *realloc_unsafe(void *old_ptr, size_t size);
static void free_unsafe(void *ptr);

static inline int umem_create(umem_head_t *ptr, umem_head_t size);
static inline void umem_remove_head(umem_head_t *prev, umem_head_t *ptr);
static inline void umem_remove_body(umem_head_t *ptr);

static umem_head_t *free_ptr;
static struct mutex_t heap_lock;

int
user_heap_init(void)
{
	 mutex_init(&heap_lock);
	 
     int err;

     if ((err = alloc_page((void *)USPACE_START, PTE_U | PTE_W)) != 0)
		  return err;

     if ((err = alloc_page((void *)USPACE_END, PTE_U | PTE_W)) != 0)
		  return err;

     free_ptr = (void *)USPACE_START;
     *free_ptr = (USPACE_END - USPACE_START) / sizeof(umem_head_t);

     // kprintf(" usable address space: %d KB\n", *free_ptr >> 8);

     /* sentinel */
     *(umem_head_t *)USPACE_END = 0;

     return 0;
}

// user heap implementation
//
// 地址空间被划完整划分成不相交的若干内存块,
// 每个内存块开始四个字节表示该内存块的(大小 / 4)
//

void *
malloc_unsafe(size_t size)
{
     /* round up the size */
#ifdef U_HEAP_DEBUG
     kprintf(" [malloc] allocating %d\n", size);
#endif
	 
     umem_head_t bsize;
     bsize = ((size + 3) >> 2) + 1;

     umem_head_t *ptr;
     ptr = free_ptr;

     while (*ptr < bsize)
     {
#ifdef U_HEAP_DEBUG
		  kprintf(" [malloc] %08x current size : %x\n", ptr, *ptr);
#endif
		  /* 合并可能的空闲内存块 */
		  while (*(ptr + *ptr) > 0)
		  {
			   umem_head_t tmp = *(ptr + *ptr);
			   umem_remove_head(ptr, ptr + *ptr);
			   *ptr += tmp;
		  }

		  if (*ptr < bsize)
		  {
			   // 连续移动两个块
			   ptr += *ptr;
			   ptr -= *ptr;

			   /* memory is not enough */
			   if (ptr >= (umem_head_t *)USPACE_END)
					return NULL;
		  }
		  else break;
     }

#ifdef U_HEAP_DEBUG
     kprintf(" [malloc] split mem (%08x, %x)\n", ptr, *ptr);
#endif
    
     umem_head_t *result;
     result = ptr;
     ptr = ptr + bsize;

     if (umem_create(result, bsize) != 0)
		  return NULL;
     
     *ptr = *result - bsize;
     *result = -bsize;

	 free_ptr = ptr;

#ifdef U_HEAP_DEBUG
     kprintf(" into mem (%08x, %x), (%08x, %x)\n", ptr, *ptr, result, -*result);
#endif
     return result + 1;
}

void *
realloc_unsafe(void *old_ptr, size_t size)
{
	 // round up the size
     /* umem_head_t bsize; */
     /* bsize = ((size + 3) >> 2) + 1; */

	 /* umem_head_t *ptr = old_ptr; */
	 
/* 	 if (-*(-- ptr) > bsize) */
/* 	 { */
/* 		  /\* 缩小 *\/ */
		  
/* #ifdef U_HEAP_DEBUG */
/* 		  kprintf(" [malloc] split mem (%08x, %x)\n", ptr, *ptr); */
/* #endif */
    
/* 		  umem_head_t *result; */
/* 		  result = ptr; */
/* 		  ptr = ptr + bsize; */

/* 		  *ptr = -*result - bsize; */
/* 		  *result = -bsize; */

/* 		  umem_remove_body(ptr); */
		  
/* #ifdef U_HEAP_DEBUG */
/* 		  kprintf(" into mem (%08x, %x), (%08x, %x)\n", result, -*result, ptr, *ptr); */
/* #endif */
/* 		  return result; */
/* 	 } */
/* 	 else */
/* 	 { */

	 /* TODO 更好的解决办法 */

	 umem_head_t *result = malloc_unsafe(size);
	 
	 if (result != NULL)
	 {
		  size_t old_size = (-*((umem_head_t *)old_ptr - 1) - 1) << 2;
		  if (size > old_size) size = old_size;
		  
		  memmove(result, old_ptr, size);
	 }
	 
	 free_unsafe(old_ptr);
	 return result;
}

void
free_unsafe(void *ptr)
{
     umem_head_t *head;
     head = ptr;
     --head;
      
     *head = -*head;
     umem_remove_body(head);

     if (free_ptr > head) free_ptr = head;
}

void
umem_remove_head(umem_head_t *prev, umem_head_t *ptr)
{
     /* assume that *ptr > 0 */
     uint32_t head_pagenum = (uint32_t)ptr >> PAGE_SHIFT;
     if ((uint32_t)prev >> PAGE_SHIFT != head_pagenum &&
		 (uint32_t)(ptr + *ptr) >> PAGE_SHIFT != head_pagenum)
     {
#ifdef U_HEAP_DEBUG
		  kprintf(" [umem_remove_head] remove %08x\n", ptr);
#endif
		  remove_map(ptr);
     }
}

void
umem_remove_body(umem_head_t *ptr)
{
     /* assume that *ptr > 0 */
     char *addr, *end_addr;
     addr = PAGE_ALIGN(ptr + 1);
     end_addr = PAGE_ALIGN(ptr + *ptr);

     if (addr == PAGE_ALIGN(ptr)) addr += PAGE_SIZE;

#ifdef U_HEAP_DEBUG
     kprintf(" [umem_remove_body] remove map [ %08x, %08x )", addr, end_addr);
#endif
     while (addr < end_addr)
     {
		  remove_map(addr);
		  addr += PAGE_SIZE;
     }
}

int
umem_create(umem_head_t *ptr, umem_head_t size)
{
     int err;
     
     char *start, *end;

     start = PAGE_ALIGN(ptr + 1);
     end   = PAGE_ALIGN(ptr + size);
     ptr   = PAGE_ALIGN(ptr);

     if ((umem_head_t *)start == ptr) start += PAGE_SIZE;
     if ((umem_head_t *)end   == PAGE_ALIGN(ptr + *ptr)) end -=PAGE_SIZE;
     
#ifdef U_HEAP_DEBUG
     kprintf(" [umem_create] create umem block, map page %08x - %08x\n", start, end);
#endif
     
     while (start <= end)
     {
		  if ((err = alloc_page(start, PTE_W | PTE_U)) != 0)
		  {
			   start -= PAGE_SIZE;
			   while ((umem_head_t *)start != ptr)
			   {
					remove_map(start);
					start -= PAGE_SIZE;
			   }
	       
			   return err;
		  }
		  start += PAGE_SIZE;
     }
     
     return 0;
}

void *
malloc(size_t size)
{
	 mutex_lock(&heap_lock);
	 void *result = malloc_unsafe(size);
	 mutex_unlock(&heap_lock);

	 return result;
}

void *
realloc(void *old_ptr, size_t size)
{
	 mutex_lock(&heap_lock);
	 void *result = realloc_unsafe(old_ptr, size);
	 mutex_unlock(&heap_lock);

	 return result;
}

void
free(void *ptr)
{
	 mutex_lock(&heap_lock);
	 free_unsafe(ptr);
	 mutex_unlock(&heap_lock);
}

/* page space management ------------------------------------------------------------ */

extern uint32_t ptab[];

int
user_page_init(void)
{
	 int err;
	 
	 if ((err = alloc_page((void *)UPAGE, PTE_U | PTE_W)) != 0)
		  return err;

	 int i;
	 for (i = 0; i != PAGETAB_COUNT; ++ i)
	 {
		  ptab[i] = i + 1;
	 }
	 ptab[PAGETAB_COUNT - 1] = 0;

	 return 0;
}

void *
palloc(void)
{
	 /* 分配失败 */
	 if (ptab[0] == 0) return NULL;
	 
	 int result = ptab[0];
	 
	 // kprintf("[palloc] %d: result = %d\n", self_idx, result);
	 
	 ptab[0] = ptab[result];

	 return (void*)(UPAGE + (result << PAGE_SHIFT));
}

void
pfree(void *addr)
{
	 int idx = ((uint32_t)addr - UPAGE) >> PAGE_SHIFT;
	 
	 ptab[idx] = ptab[0];
	 ptab[0] = idx;
}

